package com.bat.conf;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerClient;
import org.springframework.http.HttpHeaders;
import org.springframework.web.server.ServerWebExchange;

import java.net.URI;
import java.util.List;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;

/**
 * @program: gateway
 * @description:
 * @author: zhoupengfei
 * @create: 2019-08-16 18:05
 **/
@Slf4j
public class UserLoadBalancerClientFilter extends LoadBalancerClientFilter {

    public  static  ThreadLocal<ServerWebExchange> exchange=new ThreadLocal<>();
    public UserLoadBalancerClientFilter(LoadBalancerClient loadBalancer, LoadBalancerProperties properties) {
        super(loadBalancer, properties);
    }

    @Override
    protected ServiceInstance choose(ServerWebExchange exchange) {
        //这里可以拿到web请求的上下文，可以从header中取出来自己定义的数据。
        String userId = exchange.getRequest().getHeaders().getFirst("userId");
        UserLoadBalancerClientFilter.exchange.set(exchange);
        log.info("步骤1");
        if(IpUserHashRule.originHost.get()==null){
            //IpUserHashRule.originHostHeader
            HttpHeaders httpHeaders=exchange.getRequest().getHeaders();
            List<String> originHostHeader=httpHeaders.get(IpUserHashRule.originHostHeader);
            if(originHostHeader==null || originHostHeader.size()==0 ){
                String host=exchange.getRequest().getURI().getHost();
                //设置请求头
                //httpHeaders.set(IpUserHashRule.originHostHeader,host);
                exchange.getRequest().mutate().header(IpUserHashRule.originHostHeader,host).build();
                IpUserHashRule.originHost.set(host);
            }else{
                IpUserHashRule.originHost.set(originHostHeader.get(0));
            }
        }
        if (this.loadBalancer instanceof RibbonLoadBalancerClient) {
            RibbonLoadBalancerClient client = (RibbonLoadBalancerClient) this.loadBalancer;
            String serviceId = ((URI) exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).getHost();
            //这里使用userId做为选择服务实例的key
            return client.choose(serviceId, userId);
        }
        return super.choose(exchange);
    }

}